from pymodbus.client.sync import ModbusSerialClient as ModbusClient
from pymodbus.constants import Endian
from pymodbus.payload import BinaryPayloadDecoder, BinaryPayloadBuilder
import time
import math

class MecanumWheelController:
    def __init__(self, port='/dev/ttyACM0', baudrate=115200, parity='N', stopbits=1, bytesize=8, timeout=10, unit_id=0x32):
        self.client = ModbusClient(
            method='rtu',
            port=port,
            baudrate=baudrate,
            parity=parity,
            stopbits=stopbits,
            bytesize=bytesize,
            timeout=timeout
        )
        self.unit_id = unit_id

    def connect(self):
        # Connect to Modbus RTU device
        connection = self.client.connect()
        if not connection:
            print("Unable to connect to the Modbus RTU device")
            return False
        return True

    def disconnect(self):
        # Disconnect from the Modbus RTU device
        self.client.close()

    def read_real_speed(self, wheel_number=None):
        try:
            speeds = []
            if wheel_number:
                wheel_numbers = [wheel_number]
            else:
                wheel_numbers = [1, 2, 3, 4]  # Read speeds of all four wheels

            for num in wheel_numbers:
                address = 0x0006 + (num - 1)  # Calculate register address based on wheel number
                count = 1  # Read a single register
                response = self.client.read_input_registers(address, count, unit=self.unit_id)
                # time.sleep(0.01)
                if not response.isError():
                    # Decode register data
                    decoder = BinaryPayloadDecoder.fromRegisters(response.registers, byteorder=Endian.Big)
                    value = decoder.decode_16bit_uint()
                    # Convert uint16 to int16
                    value = value - 65536 if value > 32767 else value
                    
                    value = value * 1200 // (56 * 11)  # Convert encoder value to RPM
                    speeds.append(value)
                    # print(f"Wheel {num} speed: {value}")
            return speeds
        except Exception as e:
            print(f"Failed to read speed of wheel {num}: {response}")
            speeds.append(None)
            return None
        
    def set_info(self):
        try:
            # Set other information for the device
            address = 0x1F
            values = [0x0B38,0x0050,0x00C8,0x00C5]  # List of values to write
            builder = BinaryPayloadBuilder(byteorder=Endian.Big)
            for value in values:
                builder.add_16bit_uint(value)

            payload = builder.build()
            response = self.client.write_registers(address, payload, skip_encode=True, unit=self.unit_id)
            # time.sleep(0.01)
            if not response.isError():
                # print(f"Successfully wrote multiple holding registers: {values}")
                pass
        except Exception as e:
            print(f"Failed to write multiple holding registers: {e}")

    def set_speed(self, wheel_direction, wheel_number, wheel_rpm):
        try:
            # Set the running speed of a specific wheel or all wheels
            def write_speed(wheel_num):
                address = 0x0006 + (wheel_num - 1)  # Calculate register address for the wheel
                response = self.client.write_register(address, value, unit=self.unit_id)
                time.sleep(0.01)
                if not response.isError():
                    pass  # Write successful, do something if needed
            encoder_value = wheel_rpm *  (56 * 11) // 1200  # Convert RPM to encoder value
            if wheel_direction == 0:  # Counterclockwise
                value = encoder_value
            else:  # Clockwise
                value = (1 << 8) | encoder_value  # Set direction bit in high byte

            # Ensure encoder value is within 16-bit unsigned integer range
            if value < 0 or value > 0xFFFF:
                print(f"Speed value for wheel {wheel_number} is out of range: {value}")
                return

            if wheel_number == 0:
                # If wheel_number is 0, set all wheels
                for wheel_num in range(1, 5):  # Assuming TOTAL_WHEELS is the total number of wheels
                    write_speed(wheel_num)
            else:
                # Otherwise, set the specified wheel only
                write_speed(wheel_number)
        except Exception as e:
            print(f"Failed to set wheel speed: {e}")

    # Set the running speed of the device by writing to the device's holding registers.
    def set_all_speed(self,FL_speed,FR_speed,RL_speed,RR_speed):
        try:
            # Write multiple holding registers (example: writing two registers starting at address 0x02)
            address = 0x0006 # Address for motor_one
            values = [FL_speed,FR_speed,RL_speed,RR_speed] # rpm
            # Build payload
            builder = BinaryPayloadBuilder(byteorder=Endian.Big)
            for value in values:
                builder.add_16bit_uint(value)
            
            payload = builder.build()
            response = self.client.write_registers(address, payload, skip_encode=True, unit=self.unit_id)
            time.sleep(0.01) 
            if not response.isError():
                # print(f"Successfully wrote multiple holding registers: {values}")
                pass
        except Exception as e:
            print(f"Failed to write multiple holding registers: {e}")

    def set_all_staus(self, directions,speeds):
        try:
            address = 0x0006  # Assuming all motor registers are consecutive
            speeds = speeds *  (56 * 11) // 1200  
            speeds = [speeds]  * 4 # rpm
            if directions == 1:
                directions = [1, 0, 0, 1]  # Forward
            elif directions == 2:
                directions = [0, 1, 1, 0]  # Backward
            elif directions == 3:
                directions = [0, 0, 1, 1]  # Left
            elif directions == 4:
                directions = [1, 1, 0, 0]  # Right
            # Build payload
            builder = BinaryPayloadBuilder(byteorder=Endian.Big)
            for speed, direction in zip(speeds, directions):
                # Convert to hex and set direction bit
                # Assume speed value is already within valid range (less than 0x7FFF to avoid conflict with sign bit)
                if speed < 0 or speed > 0x7FFF:
                    # print(f"Speed value out of range: {speed}")
                    return  # Optionally throw an exception or continue with other processing
                value = (direction << 15) | speed  # Set direction bit in high order
                builder.add_16bit_uint(value)
            
            payload = builder.build()
            response = self.client.write_registers(address, payload, skip_encode=True, unit=self.unit_id)
            time.sleep(0.01)  # Delay time, adjust according to actual situation
            if not response.isError():
                # Write successful
                pass
        except Exception as e:
            print(f"Failed to write multiple holding registers: {e}")

    def stop(self, wheel_number=None):
        try:
            # Stop a specific wheel or all wheels
            if wheel_number != 0:
                # Stop the specified wheel number
                address = 0x0006 + (wheel_number - 1)  # Calculate register address for the wheel
                response = self.client.write_register(address, 0x0000, unit=self.unit_id)
                time.sleep(0.01)
                if not response.isError():
                    # print(f"Wheel {wheel_number} stopped successfully")
                    pass
            elif wheel_number == 0:
                # Stop all wheels
                address = 0x0006
                values = [0x0000, 0x0000, 0x0000, 0x0000]  # Set all wheel speeds to 0
                builder = BinaryPayloadBuilder(byteorder=Endian.Big)
                for value in values:
                    builder.add_16bit_uint(value)
                payload = builder.build()
                response = self.client.write_registers(address, payload, skip_encode=True, unit=self.unit_id)
                time.sleep(0.01) 
                if not response.isError():
                    # print("All wheels stopped successfully")
                    pass
        except Exception as e:
            print(f"Failed to stop wheels: {e}")

    def enable_pid(self, enable):
        try:
            # Enable or disable PID control
            address = 0x001E  # Assuming the register to enable/disable PID is at the end
            value = 1 if enable else 0  # 1 to enable PID, 0 to disable
            response = self.client.write_register(address, value, unit=self.unit_id)
            time.sleep(0.01)
            if not response.isError():
                # print(f"PID {'enabled' if enable else 'disabled'} successfully")
                pass
        except Exception as e:
            print(f"PID {'enable' if enable else 'disable'} failed: {e}")

    def set_pid_values(self, kp, ki, kd):
        try:
            # Set PID parameter values for Kp, Ki, and Kd
            address = 0x0012  # Assuming PID parameters start at register address 0x0012
            values = [kp, ki, kd] * 4  # Assuming there are four sets of PID parameters, each set to the same value
            builder = BinaryPayloadBuilder(byteorder=Endian.Big)
            for value in values:
                builder.add_16bit_uint(value)
            
            payload = builder.build()
            response = self.client.write_registers(address, payload, skip_encode=True, unit=self.unit_id)
            time.sleep(0.01) 
            if not response.isError():
                # print(f"PID parameters written successfully: Kp={kp}, Ki={ki}, Kd={kd}")
                pass
        except Exception as e:
            print(f"Failed to write PID parameters: {e}")

    def set_mode(self, mode):
        try:
            # Set the operating mode of the car
            address = 0x0A
            response = self.client.write_register(address, mode, unit=self.unit_id)
            time.sleep(0.01)
            if not response.isError():
                # print(f"Successfully wrote holding register: {mode}")
                pass
        except Exception as e:
            print(f"Failed to write holding register: {e}")

    # Convert to custom hex format
    def to_custom_hex(self,value):
        if value >= 0:
            return f"0x{value:04X}"
        else:
            return f"0x{(0x0100 + abs(value)):04X}"

    # Control the car's movement direction, linear speed, and angular speed
    def run(self,direction,speed,omega_deg):
        max_rpm = 205  # Maximum RPM value
        if direction == 0 and speed == 0:
            # print("speed")
            actual_rpm = max_rpm * omega_deg // 100
            rpm_wheel1 = actual_rpm
            rpm_wheel2 = actual_rpm
            rpm_wheel3 = actual_rpm
            rpm_wheel4 = actual_rpm
            # Get encoder value for 50ms in integers
            encoder_50ms_wheel1 = rpm_wheel1*(56*11)/1200
            encoder_50ms_wheel2 = rpm_wheel2*(56*11)/1200
            encoder_50ms_wheel3 = rpm_wheel3*(56*11)/1200
            encoder_50ms_wheel4 = rpm_wheel4*(56*11)/1200
            # Round to integers
            encoder_50ms_wheel1 = int(encoder_50ms_wheel1)
            encoder_50ms_wheel2 = int(encoder_50ms_wheel2)
            encoder_50ms_wheel3 = int(encoder_50ms_wheel3)
            encoder_50ms_wheel4 = int(encoder_50ms_wheel4)
            # Convert to hex
            ox_wheel1 = self.to_custom_hex(encoder_50ms_wheel1)
            ox_wheel2 = self.to_custom_hex(encoder_50ms_wheel2)
            ox_wheel3 = self.to_custom_hex(encoder_50ms_wheel3)
            ox_wheel4 = self.to_custom_hex(encoder_50ms_wheel4)
            # Convert string to integer
            ox_wheel1 = int(ox_wheel1, 16)
            ox_wheel2 = int(ox_wheel2, 16)
            ox_wheel3 = int(ox_wheel3, 16)
            ox_wheel4 = int(ox_wheel4, 16)
        else:
            speed = speed * 85 // 100
            omega = math.radians(omega_deg)
            vy = speed*math.cos(math.radians(-direction))
            vx = speed*math.sin(math.radians(-direction))
            omega_rad_per_s_wheel1 = (vy-vx+20*omega)/4
            omega_rad_per_s_wheel2 = (vy+vx-20*omega)/4
            omega_rad_per_s_wheel3 = (vy-vx-20*omega)/4
            omega_rad_per_s_wheel4 = (vy+vx+20*omega)/4
            # Get rpm for each wheel
            rpm_wheel1 = -(omega_rad_per_s_wheel1*60/(2*math.pi)) 
            rpm_wheel2 = omega_rad_per_s_wheel2*60/(2*math.pi)
            rpm_wheel3 = omega_rad_per_s_wheel3*60/(2*math.pi)
            rpm_wheel4 = -(omega_rad_per_s_wheel4*60/(2*math.pi)) 
            # Get encoder value for 50ms
            encoder_50ms_wheel1 = rpm_wheel1*(56*11)/1200
            encoder_50ms_wheel2 = rpm_wheel2*(56*11)/1200
            encoder_50ms_wheel3 = rpm_wheel3*(56*11)/1200
            encoder_50ms_wheel4 = rpm_wheel4*(56*11)/1200
            # Round to integers
            encoder_50ms_wheel1 = int(encoder_50ms_wheel1)
            encoder_50ms_wheel2 = int(encoder_50ms_wheel2)
            encoder_50ms_wheel3 = int(encoder_50ms_wheel3)
            encoder_50ms_wheel4 = int(encoder_50ms_wheel4)
            # Convert to hex
            ox_wheel1 = self.to_custom_hex(encoder_50ms_wheel1)
            ox_wheel2 = self.to_custom_hex(encoder_50ms_wheel2)
            ox_wheel3 = self.to_custom_hex(encoder_50ms_wheel3)
            ox_wheel4 = self.to_custom_hex(encoder_50ms_wheel4)
            # Convert string to integer
            ox_wheel1 = int(ox_wheel1, 16)
            ox_wheel2 = int(ox_wheel2, 16)
            ox_wheel3 = int(ox_wheel3, 16)
            ox_wheel4 = int(ox_wheel4, 16)
        self.set_all_speed(ox_wheel1,ox_wheel2,ox_wheel3,ox_wheel4)
        
    def set_pwm(self,direction,wheel_num,speed):
        max_rpm = 205  # Maximum RPM value
        actual_rpm = max_rpm * speed // 100     
        self.set_speed(direction,wheel_num,actual_rpm)

    # Set servo interfaces S1-S6 addresses: 0x000B, 0x000C, 0x000D, 0x000E, 0x000F, 0x0010
    def set_degree(self, address, degree):
        try:
            # Set the angle for servo interfaces S1-S6
            response = self.client.write_register(address, degree, unit=self.unit_id)
            time.sleep(0.01)
            if not response.isError():
                # print(f"Successfully wrote holding register: {degree}")
                pass
        except Exception as e:
            print(f"Failed to write holding register: {e}")
